package com.dragon.aipictures.test2;

import org.apache.curator.shaded.com.google.common.hash.Hashing;

import java.util.Collection;
import java.util.SortedMap;
import java.util.TreeMap;

// 分库算法：一致性哈希分片
public class DbConsistentHashAlgorithm {
    private final TreeMap<Long, String> virtualNodes = new TreeMap<>();
    private static final int VIRTUAL_NODES_PER_DB = 1000;

    public void init(Collection<String> databases) {
        for (String db : databases) {
            for (int i=0; i<VIRTUAL_NODES_PER_DB; i++) {
                long hash = murmurHash("DB-" + db + "-VN-" + i);
                virtualNodes.put(hash, db);
            }
        }
    }

    public String route(String shardKey) {
        long hash = murmurHash(shardKey);
        SortedMap<Long, String> tailMap = virtualNodes.tailMap(hash);
        return tailMap.isEmpty() ? virtualNodes.firstEntry().getValue() : tailMap.get(tailMap.firstKey());
    }
    private long murmurHash(String key) {
        return Hashing.murmur3_128().hashUnencodedChars(key).asLong();
    }
}